import pandas as pd
from pulp import LpMaximize, LpProblem, LpVariable, lpSum

# 假设股票价格（单位：元/股），需要根据实际情况调整
# ['600036', '002594', '300006', '000905', '002607']
stock_prices = [36.69, 202.27, 7.69, 7.31, 34.2]

# 读取CSV文件并计算平均收益率
def calculate_avg_return(file_path):
    df = pd.read_csv(file_path)
    df['date'] = pd.to_datetime(df['date'])
    df = df[(df['date'] >= '2021-01-01') & (df['date'] <= '2022-06-30')]
    avg_return = df['return'].mean()
    return avg_return

# 假设CSV文件路径
base_path = './data_return/'
csv_files = [base_path + code + '_return.csv' for code in ['600036', '002594', '300006', '000905', '002607']]


# 计算每支股票的平均收益率
avg_returns = [calculate_avg_return(file) for file in csv_files]

# 创建整数线性规划问题
problem = LpProblem(name="stock-investment", sense=LpMaximize)

# 创建决策变量：每支股票的购买股数（单位：手）
x = {i: LpVariable(name=f"x_{i}", lowBound=20, cat='Integer') for i in range(5)}

# 添加目标函数：最大化投资收益
problem += lpSum([avg_returns[i] * x[i] * 100 for i in range(5)])

# 添加约束条件：总投资金额不超过200万元
problem += lpSum([stock_prices[i] * x[i] * 100 for i in range(5)]) <= 2000000

# 求解问题
status = problem.solve()

# 输出结果
if status == 1:  # 1表示最优解
    print("最优投资配置：")
    total_investment = 0
    for i in range(5):
        stock_num = x[i].value()
        stock_investment = stock_prices[i] * stock_num * 100
        total_investment += stock_investment
        print(f"股票{i+1}：购买股数 = {stock_num}手，投资金额 = {stock_investment:.2f}元")
    print(f"总的投资金额：{total_investment:.2f}元")
else:
    print("未找到最优解，请检查数据或约束条件。")